package main

import (
	"bufio"
	"database/sql"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
)

var wg sync.WaitGroup
var batchSize = 5000

func main() {
	self, _ := filepath.Abs(os.Args[0])
	log.Println("正在执行数据库初始化程序> ", self)

	var connStr string
	fmt.Println("\n请输入数据库连接，格式 user:pwd@tcp(host:port)")
	fmt.Scanln(&connStr)

	if isBlank(connStr) {
		log.Fatal("未获取到数据库连接信息，程序终止")
	}

	var confirm string
	fmt.Println("\n是否确认将该应用目录下*.sql文件初始化到数据库，文件名即库名称？(y/n)")
	fmt.Scanln(&confirm)

	if strings.ToUpper(confirm) != "Y" {
		os.Exit(0)
	}

	// 获取脚本文件所在目录的路径
	dir := filepath.Dir(self)

	// 获取目录下的所有 .sql 文件
	files, err := filepath.Glob(filepath.Join(dir, "*.sql"))
	checkError(err)

	// 连接到数据库
	db, err := sql.Open("mysql", connStr+"/")
	checkError(err)
	defer db.Close()

	// 遍历每个 .sql 文件并执行 SQL 语句
	for index, file := range files {
		wg.Add(1)
		log.Printf("正在处理第%d个文件：%s", index+1, file)
		go execute(file, connStr, db)
		log.Println()
	}

	wg.Wait()
}

func execute(path string, connStr string, odb *sql.DB) {
	defer wg.Done()

	//创建数据库并使用数据库
	dbName := strings.TrimSuffix(filepath.Base(path), ".sql")
	odb.Exec(fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s CHARSET utf8mb4;", dbName))

	// 打开SQL语句文件
	file, err := os.Open(path)
	checkError(err)
	defer file.Close()

	// 逐行读取 SQL 文件内容
	scanner := bufio.NewScanner(file)

	var inserts, others []string
	var query string
	for scanner.Scan() {
		line := scanner.Text()
		// 忽略 -- 开头
		if strings.HasPrefix(line, "--") {
			continue
		}
		query += line

		// 如果以分号结尾，则执行该语句
		if len(line) > 0 && line[len(line)-1:] == ";" && !isBlank(query) {
			if strings.HasPrefix(query, "INSERT") || strings.HasPrefix(query, "UPDATE") || strings.HasPrefix(query, "DELETE") {
				inserts = append(inserts, strings.TrimSpace(query))
			} else {
				others = append(others, strings.TrimSpace(query))
			}

			// 清空 query
			query = ""
		}
	}

	//为协程创建新的数据库连接
	db, err := sql.Open("mysql", connStr+"/"+dbName)
	//db.SetMaxOpenConns(20)
	//db.SetMaxIdleConns(20)
	defer db.Close()

	tx, err := db.Begin()

	for _, query := range others {
		tx.Exec(query)
	}

	// 关闭外键检查 ，同一个数据库连接中，只要设置一次就够了
	sqlCloseFKC := `SET FOREIGN_KEY_CHECKS = 0;`
	_, _ = tx.Exec(sqlCloseFKC)

	for i, query := range inserts {
		tx.Exec(query)

		if i%batchSize == 0 {
			_ = tx.Commit()
			tx, _ = db.Begin()
		}

	}

	_ = tx.Commit()

	log.Println("数据库：", dbName, " 初始化完毕")
}

func checkError(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

func isBlank(s string) bool {
	if len(s) == 0 {
		return true
	}
	r := regexp.MustCompile(`^\s*$`) // 匹配只包含空白字符的正则表达式
	return r.MatchString(s)
}
